# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/4 上午11:47
@Author yzpang
@Function: 
"""
import torch

from modelserver.configs.model_config import TOP_K


def text_classification(model, tokenizer, text):
    """
    文本分类
    :param tokenizer: 分词器
    :param model: 模型
    :param text: 输入文本 e.g.  '金属卡扣安装不到位，上海乐扣乐扣贸易有限公司将召回捣碎器1162件'
    :return: 得分最高的前几条 e.g. ([9.730793952941895, 1.1261529922485352, 0.9946187734603882, 0.8857908248901367, 0.835823655128479], [8, 6, 61, 37, 5])
    """
    # next(model.parameters()).device 模型使用的设备
    inputs = tokenizer(text, return_tensors='pt').to(next(model.parameters()).device)
    with torch.no_grad():  # 不需要计算梯度
        outputs = model(**inputs)

    # 获取logits并找到分数最高的类别
    logits = outputs.logits.cpu()
    # 获取logits并找到分数最高的5个类别
    top_scores, top_indices = torch.topk(logits, k=TOP_K, dim=-1)
    top_scores = top_scores.squeeze().tolist()
    top_indices = top_indices.squeeze().tolist()
    return top_scores, top_indices
